import torch
from torchvision import transforms, models
from PIL import Image
import os

# -----------------------------
# 🟢 配置参数
# -----------------------------
MODEL_PATH = "cat_dog_model_mps.pth"  # 已训练模型路径
DEVICE = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
CLASS_NAMES = ["cat", "dog"]           # 类别顺序需与训练时一致
CONFIDENCE_THRESHOLD = 0.7             # 低于阈值判为未知类别

# -----------------------------
# 🟢 加载模型
# -----------------------------
model = models.resnet18(weights=None)
model.fc = torch.nn.Linear(model.fc.in_features, len(CLASS_NAMES))
model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
model.to(DEVICE)
model.eval()

print(f"📌 当前设备: {DEVICE}")
print(f"✅ 模型已成功加载: {MODEL_PATH}")

# -----------------------------
# 🟢 图片预处理函数
# -----------------------------
def preprocess_image(image_path):
    transform = transforms.Compose([
        transforms.Resize((128, 128)),
        transforms.ToTensor(),
    ])
    image = Image.open(image_path).convert("RGB")
    return transform(image).unsqueeze(0)  # 增加 batch 维度

# -----------------------------
# 🟢 单张图片推理函数
# -----------------------------
def predict(image_path):
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"图片不存在: {image_path}")

    input_tensor = preprocess_image(image_path).to(DEVICE)
    with torch.no_grad():
        logits = model(input_tensor)
        print("🧮 logits:", logits)

        # Softmax 计算概率
        probs = torch.softmax(logits, dim=1)
        probs_list = probs[0].tolist()
        softmax_dict = {cls: round(p, 4) for cls, p in zip(CLASS_NAMES, probs_list)}
        print("📊 Softmax分布:", softmax_dict)

        # Logits 差距
        logits_sorted, _ = torch.sort(logits, dim=1, descending=True)
        if logits_sorted.shape[1] >= 2:
            logits_diff = (logits_sorted[0][0] - logits_sorted[0][1]).item()
        else:
            logits_diff = 0.0
        print(f"📉 Logits差距: {logits_diff:.4f}")

        # 预测类别
        pred_idx = torch.argmax(probs, dim=1).item()
        pred_class = CLASS_NAMES[pred_idx]
        pred_prob = probs[0][pred_idx].item()

        # 置信度阈值判断
        if pred_prob < CONFIDENCE_THRESHOLD:
            pred_class = "未知类别"

    return pred_class, pred_prob, softmax_dict, logits_diff

# -----------------------------
# 🟢 测试单张图片
# -----------------------------
if __name__ == "__main__":
    test_image_path = "/Users/up_dong/Documents/python_workspace/cats-vs-dogs/data/test/img.png"

    try:
        pred_class, pred_prob, softmax_probs, logits_diff = predict(test_image_path)
    except FileNotFoundError:
        print(f"❌ 图片不存在: {test_image_path}")
        exit(1)

    print("\n✅ 预测结果:")
    print(f"🎯 类别: {pred_class}")
    print(f"📊 置信度: {pred_prob:.4f}")
    print(f"📁 图片路径: {test_image_path}")
